import numpy as np
from tables import *

def createPhaseSpaceOutDB(numBunches):
    PhaseSpaceOut = {
        "dE" : Float32Col(shape=numBunches),
        "dt" : Float32Col(shape=numBunches),
        "emittance": Float32Col(),
        "rmsdE": Float32Col(),
        "rmsdt": Float32Col(),
        }
    return PhaseSpaceOut

def createCavityPhaseSpace(numCavs):
    CavityPhaseSpace = {
            "Index" : Int64Col(),
            "dE" : Float32Col(shape=numCavs),
            "dt" : Float32Col(shape=numCavs),
            }
    return CavityPhaseSpace

def createLinacInfoDB(self,numBunches,numCavs):
    LinacInfo = {
        "Pulse" : {
        "numBunches" : Int64Col(),
        "numCavs" : Int64Col(),
        "tb" : Float64Col(),
        "Ib" : Float64Col(),
        "dq" : Float64Col(shape=numBunches),
        "dEIn" : Float64Col(shape=(numBunches)),
        "dtIn" : Float64Col(shape=(numBunches)),
        "dEOut" : Float64Col(shape=(numBunches)),
        "dtOut" : Float64Col(shape=(numBunches)),
             },
         "HOM" : {
            "freq" : Float64Col(shape=numCavs),
            "Qex" : Float64Col(shape=numCavs),
            "RoverQ" : Float64Col(shape=numCavs),
            "HOMV" : Float64Col(shape=numCavs),
            "Vind" : Float64Col(shape=numCavs),
            },
 
    }
    return LinacInfo

class LinacDB:
    def __init__(self,filename,numBunches,mode='r'):
        self.h5file = openFile(filename,mode)
        self.root = self.h5file.root
        self.filters = Filters(complevel=1, complib='blosc')

    def AddLinac(self,name,numBunches,numCavs):
        group = self.h5file.createGroup(self.root,name,title)
        LInfo = createPhaseOutDB(numBunches,numCavs)
        self.Infotable = self.h5file.createTable("/"+name, 'Info', LInfo,expectedrows=1)#,filters=self.filters)

    def Write(self,LinObj,table,Full=False):
        L = table.row 
        L['emittance']=LinObj.Pulse.emittance
        L['rmsdE']=LinObj.Pulse.rmsdE
        L['rmsdt']=LinObj.Pulse.rmsdt
        if Full==True:
            L['dt'] = LinObj.Pulse.dt[:,-1]
            L['dE'] = LinObj.Pulse.dE[:,-1]
        L.append()
        table.flush()

    def WriteLossInfo(self,LinObj,table,Full=False):
        ind = np.flatnonzero(LinObj.Pulse.Lost==0)
        L = table.row 
        for i in ind:
            L['Index'] = i
            if Full==True:
                L['dt'] = LinObj.Pulse.dt[i,:]
                L['dE'] = LinObj.Pulse.dE[i,:]
            L.append()
        table.flush()

    def WriteLong(self,LinObj):
        L = self.PStable.row
        L['dt'] = LinObj.Pulse.dt[:,-1]
        L['dE'] = LinObj.Pulse.dE[:,-1]
        L['emittance']=LinObj.Pulse.emittance
        L.append()
        if self.Infotable:
            L = self.Infotable.row
            L['Pulse/numBunches'] = LinObj.Pulse.numBunches
            L['Pulse/numCavs'] = LinObj.Cavs.cavNum
            L['Pulse/tb'] = LinObj.Pulse.tb
            L['Pulse/Ib'] = LinObj.Pulse.Ib
            L['Pulse/dq'] = LinObj.Pulse.dq
            L['Pulse/dEIn'] = LinObj.Pulse.dE[:,0]
            L['Pulse/dtIn'] = LinObj.Pulse.dt[:,0]
            L['Pulse/dtOut'] = LinObj.Pulse.dt[:,-1]
            L['Pulse/dtOut'] = LinObj.Pulse.dt[:,-1]
            if LinObj.HOM:
                L['HOM/freq'] = LinObj.HOM.freq
                L['HOM/Qex'] = LinObj.HOM.Qex
                L['HOM/RoverQ'] = LinObj.HOM.RoverQ
                L['HOM/HOMV'] = LinObj.HOM.HOMV
                L['HOM/Vind'] = LinObj.HOM.Vind
            L.append()
        if self.Datatable:
            L = self.Datatable.row
            for i in xrange(LinObj.Cavs.cavNum+2):
                L['Pulse/dE'] = LinObj.Pulse.dE[:,i]
                L['Pulse/dt'] = LinObj.Pulse.dt[:,i]
                L.append()
        self.PStable.flush()
        if self.Infotable:
            self.Infotable.flush()
        if self.Datatable:
            self.Datatable.flush()

